import numpy as np
import csv
import os

def read_transformation_matrices_from_files(folder_path):
    """
    从指定文件夹中的txt文件读取变换矩阵，并将它们转换为一个二维numpy数组。
    
    参数:
        folder_path (str): 包含txt文件的文件夹路径。
        
    返回:
        numpy.ndarray: 一个二维数组，形状为 (n, 12)，其中n是txt文件的数量，12是每个变换矩阵的有意义参数数量。
    """
    # 获取文件夹中所有txt文件的列表
    txt_files = [f for f in os.listdir(folder_path) if f.endswith('.txt')]
    txt_files.sort()  # 按文件名排序，确保按顺序读取
    print(txt_files)
    
    matrices = []
    
    for file_name in txt_files:
        file_path = os.path.join(folder_path, file_name)
        
        # 读取txt文件内容
        with open(file_path, 'r') as file:
            lines = file.readlines()
        
        # 将每行内容转换为浮点数，并组成一个4x4的矩阵
        matrix = []
        for line in lines:
            row = [float(num) for num in line.strip().split()]
            matrix.append(row)
        
        # 转换为numpy数组
        matrix = np.array(matrix)
        
        # 提取有意义的12个参数（假设最后一个行是[0, 0, 0, 1]，所以忽略它）
        meaningful_params = matrix[:3, :].flatten()  # 取前3行的所有元素，展平为1维
        
        matrices.append(meaningful_params)
    
    # 将所有有意义的参数组合成一个二维数组
    result = np.array(matrices)
    
    return result

def save_transformation_matrices_to_csv(matrices, csv_file_path):
    """
    将变换矩阵的NumPy数组保存到CSV文件中，并添加指定的表头。
    如果文件不存在，则新建文件并写入表头；如果文件存在，则直接追加数据。
    
    参数:
        matrices (numpy.ndarray): 变换矩阵数组，形状为 (n, 12)
        csv_file_path (str): 要保存的CSV文件路径
    """
    # 定义表头
    header = ["Matrix Index", "Vx0", "Vy0", "Vz0", "P0", "Vx1", "Vy1", "Vz1", "P1", "Vx2", "Vy2", "Vz2", "P2"]
    
    # 检查文件是否存在
    file_exists = os.path.exists(csv_file_path)
    
    # 打开文件，如果文件不存在则写入表头，如果存在则直接追加数据
    with open(csv_file_path, 'a', newline='') as csvfile:
        writer = csv.writer(csvfile)
        
        # 如果文件不存在，写入表头
        if not file_exists:
            writer.writerow(header)
        
        # 为每一行数据添加Matrix Index
        for i, row in enumerate(matrices):
            writer.writerow([i+1] + row.tolist())

# 示例用法
folder_path = 'ob_in_cam'  # 替换为你的txt文件所在的文件夹路径
csv_file_path = 'foundation_pose_matrices.csv'  # 替换为你的CSV文件路径

# 读取txt文件中的变换矩阵
transformation_matrices = read_transformation_matrices_from_files(folder_path)

# 将变换矩阵保存到CSV文件
save_transformation_matrices_to_csv(transformation_matrices, csv_file_path)